#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AKShare基金排行数据下载器
功能：使用AKShare的fund_info_index_em接口下载基金排行数据并存储到MySQL数据库
作者：AI Assistant
创建时间：2025年
"""

import os
import sys
import time
import logging
import re
from datetime import datetime, date
from typing import List, Dict, Optional, Any
from decimal import Decimal

# 尝试导入Django相关模块
try:
    import django
    import pandas as pd
    
    # 设置Django环境
    current_dir = os.path.dirname(os.path.abspath(__file__))
    sys.path.insert(0, current_dir)
    os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'wealth_backend.settings')
    django.setup()
    
    from data_service.models import FundNavRanking
    from django.db import transaction
    from django.utils import timezone
except ImportError as e:
    print(f"导入Django模块失败: {e}")
    sys.exit(1)

# 配置日志
log_dir = os.path.join(current_dir, 'log')
os.makedirs(log_dir, exist_ok=True)

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# 清除现有处理器
for handler in logger.handlers[:]:
    logger.removeHandler(handler)

# 配置日志格式
file_formatter = logging.Formatter('%(asctime)s - %(filename)s:%(lineno)d - %(levelname)s - %(message)s')
console_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')

# 创建文件处理器
file_handler = logging.FileHandler(os.path.join(log_dir, 'fund_ranking_akshare.log'))
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(file_formatter)

# 创建控制台处理器
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)
console_handler.setFormatter(console_formatter)

# 添加处理器到logger
logger.addHandler(file_handler)
logger.addHandler(console_handler)

# 防止日志传播到根logger
logger.propagate = False

# 配置参数
CONFIG = {
    'REQUEST_DELAY': 1,      # 请求间隔（秒）
    'MAX_RETRIES': 3,        # 最大重试次数
    'BATCH_SIZE': 100,       # 批量保存大小
    'TIMEOUT': 30,           # 请求超时（秒）
    'RETRY_DELAY_BASE': 2,   # 重试延迟基数（指数退避）
}


class AKShareFundRankingDownloader:
    """AKShare基金排行数据下载器"""
    
    def __init__(self):
        """
        初始化AKShare基金排行数据下载器
        """
        # 检查数据库设置
        self.check_database_setup()
        # symbol参数选项
        self.symbol_options = [
            "全部", 
            "沪深指数", 
            "行业主题", 
            "大盘指数", 
            "中盘指数", 
            "小盘指数", 
            "股票指数", 
            "债券指数"
            ]
        
        # indicator参数选项
        self.indicator_options = [
            "全部",
            "被动指数型",
            "增强指数型"

        ]
        
        self.total_combinations = len(self.symbol_options) * len(self.indicator_options)
        self.current_combination = 0
        
        logger.info(f"初始化AKShare基金排行数据下载器，共{self.total_combinations}个参数组合")
    
    def check_database_setup(self):
        """
        检查数据库表是否存在，不存在则创建
        """
        try:
            from django.db import connection
            with connection.cursor() as cursor:
                cursor.execute("SHOW TABLES LIKE 'fund_nav_ranking'")
                if not cursor.fetchone():
                    logger.warning("数据库表不存在，正在创建...")
                    # 创建表的SQL语句
                    create_table_sql = """
                    CREATE TABLE IF NOT EXISTS fund_nav_ranking (
                        id BIGINT AUTO_INCREMENT PRIMARY KEY,
                        fund_code VARCHAR(20) NOT NULL,
                        fund_name VARCHAR(200) NOT NULL,
                        nav DECIMAL(10,4),
                        accumulated_nav DECIMAL(10,4),
                        growth_rate_1d DECIMAL(8,4),
                        growth_rate_1w DECIMAL(8,4),
                        growth_rate_1m DECIMAL(8,4),
                        growth_rate_3m DECIMAL(8,4),
                        growth_rate_6m DECIMAL(8,4),
                        growth_rate_1y DECIMAL(8,4),
                        growth_rate_2y DECIMAL(8,4),
                        growth_rate_3y DECIMAL(8,4),
                        growth_rate_since_inception DECIMAL(8,4),
                        nav_date DATE,
                        data_date DATE NOT NULL,
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                        UNIQUE KEY unique_fund_data_date (fund_code, data_date),
                        INDEX idx_fund_code (fund_code),
                        INDEX idx_data_date (data_date),
                        INDEX idx_nav_date (nav_date)
                    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
                    """
                    cursor.execute(create_table_sql)
                    logger.info("数据库表创建成功")
                else:
                    logger.info("数据库表已存在")
        except Exception as e:
            logger.error(f"数据库检查失败: {e}")
            raise
    
    def get_fund_info_with_retry(self, symbol: str = "全部", indicator: str = "全部", max_retries: Optional[int] = None) -> Optional[pd.DataFrame]:
        """
        带重试机制的数据获取
        
        Args:
            symbol: 基金类型
            indicator: 指标类型
            max_retries: 最大重试次数，默认使用配置值
            
        Returns:
            DataFrame或None
        """
        if max_retries is None:
            max_retries = CONFIG['MAX_RETRIES']
            
        for attempt in range(max_retries):
            try:
                df = self.get_fund_info_index_em(symbol, indicator)
                if df is not None and not df.empty:
                    return df
                else:
                    logger.warning(f"第{attempt+1}次尝试未获取到数据")
            except Exception as e:
                logger.warning(f"第{attempt+1}次尝试失败: {e}")
                if attempt < max_retries - 1:
                    delay = CONFIG['RETRY_DELAY_BASE'] ** attempt
                    logger.info(f"等待{delay}秒后重试...")
                    time.sleep(delay)
                else:
                    logger.error(f"重试{max_retries}次后仍失败")
        return None
    
    def get_fund_info_index_em(self, symbol: str = "全部", indicator: str = "全部") -> Optional[pd.DataFrame]:
        """
        调用AKShare的fund_info_index_em接口获取基金数据
        
        Args:
            symbol: 基金类型，可选值："全部", "沪深指数", "行业主题", "策略指数", "综合指数", "债券指数", "海外指数"
            indicator: 指标类型，可选值："全部", "被动指数型", "增强指数型", "ETF"
            
        Returns:
            DataFrame或None
        """
        try:
            try:
                import akshare as ak
            except ImportError:
                logger.error("akshare模块未安装，请运行: pip install akshare")
                return None
            
            logger.info(f"开始获取基金数据 - symbol: {symbol}, indicator: {indicator}")
            
            # 调用AKShare接口
            df = ak.fund_info_index_em(symbol=symbol, indicator=indicator)
            
            if df is not None and not df.empty:
                logger.info(f"成功获取{len(df)}条基金数据 - symbol: {symbol}, indicator: {indicator}")
                return df
            else:
                logger.warning(f"未获取到数据 - symbol: {symbol}, indicator: {indicator}")
                return None
                
        except Exception as e:
            logger.error(f"获取基金数据失败 - symbol: {symbol}, indicator: {indicator}, 错误: {str(e)}")
            return None
    
    def parse_fund_data(self, df: pd.DataFrame, symbol: str, indicator: str) -> List[Dict[str, Any]]:
        """
        解析基金数据DataFrame为字典列表
        
        Args:
            df: AKShare返回的DataFrame
            symbol: symbol分类
            indicator: indicator类型
            
        Returns:
            解析后的基金数据列表
        """
        fund_data_list = []
        
        try:
            for index, row in df.iterrows():
                fund_data = {
                    'fund_code': str(row.get('基金代码', '')).strip(),
                    'fund_name': str(row.get('基金名称', '')).strip(),
                    'unit_nav': self._safe_decimal(row.get('单位净值')),
                    'nav_date': self._safe_date(row.get('日期')),
                    'daily_growth_rate': self._safe_decimal(row.get('日增长率')),
                    'weekly_growth_rate': self._safe_decimal(row.get('近1周')),
                    'monthly_growth_rate': self._safe_decimal(row.get('近1月')),
                    'quarterly_growth_rate': self._safe_decimal(row.get('近3月')),
                    'half_year_growth_rate': self._safe_decimal(row.get('近6月')),
                    'yearly_growth_rate': self._safe_decimal(row.get('近1年')),
                    'two_year_growth_rate': self._safe_decimal(row.get('近2年')),
                    'three_year_growth_rate': self._safe_decimal(row.get('近3年')),
                    'ytd_growth_rate': self._safe_decimal(row.get('今年来')),
                    'since_inception_growth_rate': self._safe_decimal(row.get('成立来')),
                    'fee_rate': self._safe_decimal(row.get('手续费')),
                    'min_purchase_amount': str(row.get('起购金额', '')).strip(),
                    'tracking_target': str(row.get('跟踪标的', '')).strip(),
                    'tracking_method': str(row.get('跟踪方式', '')).strip(),
                    'symbol_category': symbol,
                    'indicator_type': indicator
                }
                
                # 验证数据有效性后添加
                if self.validate_fund_data(fund_data):
                    fund_data_list.append(fund_data)
            
            logger.info(f"解析完成，有效数据{len(fund_data_list)}条")
            
        except Exception as e:
            logger.error(f"解析基金数据失败: {str(e)}")
        
        return fund_data_list
    
    def log_progress(self, current: int, total: int, records: int):
        """
        记录进度信息
        
        Args:
            current: 当前进度
            total: 总数
            records: 已获取记录数
        """
        progress = (current / total) * 100 if total > 0 else 0
        logger.info(f"进度: {current}/{total} ({progress:.1f}%), 已获取记录: {records}")
    
    def _safe_decimal(self, value) -> Optional[Decimal]:
        """
        安全转换为Decimal类型
        """
        if pd.isna(value) or value == '' or value == '-':
            return None
        try:
            # 移除百分号
            if isinstance(value, str) and '%' in value:
                value = value.replace('%', '')
            return Decimal(str(value))
        except (ValueError, TypeError, Exception):
            return None
    
    def _safe_date(self, value) -> Optional[date]:
        """
        安全转换为日期类型
        """
        if pd.isna(value) or value == '':
            return None
        try:
            if isinstance(value, str):
                # 尝试解析不同的日期格式
                for fmt in ['%Y-%m-%d', '%Y/%m/%d', '%Y.%m.%d']:
                    try:
                        return datetime.strptime(value, fmt).date()
                    except ValueError:
                        continue
            elif hasattr(value, 'date'):
                return value.date()
            return None
        except Exception:
            return None
    
    def validate_fund_data(self, fund_data: Dict[str, Any]) -> bool:
        """
        验证基金数据的完整性
        
        Args:
            fund_data: 基金数据字典
            
        Returns:
            是否有效
        """
        # 检查必填字段
        required_fields = ['fund_code', 'fund_name']
        for field in required_fields:
            if not fund_data.get(field) or str(fund_data[field]).strip() == '':
                logger.debug(f"基金数据缺少必填字段: {field}")
                return False
        
        # 验证基金代码格式（6位数字）
        fund_code = str(fund_data['fund_code']).strip()
        if not re.match(r'^\d{6}$', fund_code):
            logger.debug(f"基金代码格式异常: {fund_code}")
            return False
        
        return True
    

    
    def save_fund_data_batch(self, fund_data_list: List[Dict[str, Any]], batch_size: Optional[int] = None) -> bool:
        """
        批量保存数据，提高性能
        
        Args:
            fund_data_list: 基金数据列表
            batch_size: 批量大小，默认使用配置值
            
        Returns:
            是否保存成功
        """
        if not fund_data_list:
            logger.warning("没有基金数据需要保存")
            return True
            
        if batch_size is None:
            batch_size = CONFIG['BATCH_SIZE']
        
        try:
            total_saved = 0
            total_updated = 0
            
            with transaction.atomic():
                for i in range(0, len(fund_data_list), batch_size):
                    batch = fund_data_list[i:i + batch_size]
                    fund_objects = []
                    
                    for fund_data in batch:
                        if self.validate_fund_data(fund_data):
                            fund_objects.append(FundNavRanking(
                                 fund_code=fund_data['fund_code'],
                                 fund_name=fund_data['fund_name'],
                                 unit_nav=fund_data['unit_nav'],
                                 nav_date=fund_data['nav_date'],
                                 daily_growth_rate=fund_data['daily_growth_rate'],
                                 weekly_growth_rate=fund_data['weekly_growth_rate'],
                                 monthly_growth_rate=fund_data['monthly_growth_rate'],
                                 quarterly_growth_rate=fund_data['quarterly_growth_rate'],
                                 half_year_growth_rate=fund_data['half_year_growth_rate'],
                                 yearly_growth_rate=fund_data['yearly_growth_rate'],
                                 two_year_growth_rate=fund_data['two_year_growth_rate'],
                                 three_year_growth_rate=fund_data['three_year_growth_rate'],
                                 ytd_growth_rate=fund_data.get('ytd_growth_rate'),
                                 since_inception_growth_rate=fund_data['since_inception_growth_rate'],
                                 fee_rate=fund_data.get('fee_rate'),
                                 min_purchase_amount=fund_data.get('min_purchase_amount'),
                                 tracking_target=fund_data.get('tracking_target'),
                                 tracking_method=fund_data.get('tracking_method'),
                                 symbol_category=fund_data.get('symbol_category'),
                                 indicator_type=fund_data.get('indicator_type')
                             ))
                    
                    if fund_objects:
                        # 使用逐条保存确保正确处理重复数据
                        try:
                            saved, updated = self._save_individual_records(batch)
                            total_saved += saved
                            total_updated += updated
                            logger.info(f"批量处理完成 - 新增: {saved}条, 更新: {updated}条")
                        except Exception as e:
                            logger.error(f"批量处理失败: {e}")
                            continue
            
            logger.info(f"数据保存完成 - 新增: {total_saved}条, 更新: {total_updated}条")
            return True
            
        except Exception as e:
            logger.error(f"批量保存基金数据失败: {str(e)}")
            return False
    
    def _save_individual_records(self, fund_data_list: List[Dict[str, Any]]) -> tuple:
        """
        逐条保存记录（批量保存失败时的备用方案）
        
        Args:
            fund_data_list: 基金数据列表
            
        Returns:
            (saved_count, updated_count)
        """
        saved_count = 0
        updated_count = 0
        
        for fund_data in fund_data_list:
            try:
                obj, created = FundNavRanking.objects.update_or_create(
                            fund_code=fund_data['fund_code'],
                            nav_date=fund_data['nav_date'],
                            symbol_category=fund_data['symbol_category'],
                            indicator_type=fund_data['indicator_type'],
                     defaults={
                         'fund_name': fund_data['fund_name'],
                         'unit_nav': fund_data['unit_nav'],
                         'daily_growth_rate': fund_data['daily_growth_rate'],
                         'weekly_growth_rate': fund_data['weekly_growth_rate'],
                         'monthly_growth_rate': fund_data['monthly_growth_rate'],
                         'quarterly_growth_rate': fund_data['quarterly_growth_rate'],
                         'half_year_growth_rate': fund_data['half_year_growth_rate'],
                         'yearly_growth_rate': fund_data['yearly_growth_rate'],
                         'two_year_growth_rate': fund_data['two_year_growth_rate'],
                         'three_year_growth_rate': fund_data['three_year_growth_rate'],
                         'since_inception_growth_rate': fund_data['since_inception_growth_rate']
                     }
                 )
                
                if created:
                    saved_count += 1
                else:
                    updated_count += 1
                    
            except Exception as e:
                logger.error(f"保存基金数据失败 - 基金代码: {fund_data.get('fund_code')}, 错误: {str(e)}")
                continue
        
        return saved_count, updated_count
    
    def save_fund_data_to_database(self, fund_data_list: List[Dict[str, Any]]) -> bool:
        """
        将基金数据保存到数据库，使用INSERT ON DUPLICATE KEY UPDATE语法
        
        Args:
            fund_data_list: 基金数据列表
            
        Returns:
            是否保存成功
        """
        if not fund_data_list:
            logger.warning("没有基金数据需要保存")
            return True
        
        try:
            from django.db import connection
            
            with transaction.atomic():
                cursor = connection.cursor()
                
                # 构建批量插入的SQL语句
                sql = """
                INSERT INTO fund_nav_ranking (
                    fund_code, fund_name, nav_date, unit_nav, 
                    daily_growth_rate, weekly_growth_rate, monthly_growth_rate,
                    quarterly_growth_rate, half_year_growth_rate, yearly_growth_rate,
                    two_year_growth_rate, three_year_growth_rate, since_inception_growth_rate,
                    symbol_category, indicator_type, created_at, updated_at
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, NOW(), NOW())
                ON DUPLICATE KEY UPDATE
                    fund_name = VALUES(fund_name),
                    unit_nav = VALUES(unit_nav),
                    daily_growth_rate = VALUES(daily_growth_rate),
                    weekly_growth_rate = VALUES(weekly_growth_rate),
                    monthly_growth_rate = VALUES(monthly_growth_rate),
                    quarterly_growth_rate = VALUES(quarterly_growth_rate),
                    half_year_growth_rate = VALUES(half_year_growth_rate),
                    yearly_growth_rate = VALUES(yearly_growth_rate),
                    two_year_growth_rate = VALUES(two_year_growth_rate),
                    three_year_growth_rate = VALUES(three_year_growth_rate),
                    since_inception_growth_rate = VALUES(since_inception_growth_rate),
                    updated_at = NOW()
                """
                
                # 准备批量数据
                batch_data = []
                for fund_data in fund_data_list:
                    batch_data.append((
                        fund_data['fund_code'],
                        fund_data['fund_name'],
                        fund_data['nav_date'],
                        fund_data['unit_nav'],
                        fund_data['daily_growth_rate'],
                        fund_data['weekly_growth_rate'],
                        fund_data['monthly_growth_rate'],
                        fund_data['quarterly_growth_rate'],
                        fund_data['half_year_growth_rate'],
                        fund_data['yearly_growth_rate'],
                        fund_data['two_year_growth_rate'],
                        fund_data['three_year_growth_rate'],
                        fund_data['since_inception_growth_rate'],
                        fund_data['symbol_category'],
                        fund_data['indicator_type']
                    ))
                
                # 执行批量插入/更新
                cursor.executemany(sql, batch_data)
                affected_rows = cursor.rowcount
                
                logger.info(f"数据保存完成 - 影响行数: {affected_rows}条")
                return True
                
        except Exception as e:
            logger.error(f"批量保存基金数据失败: {str(e)}")
            return False
    
    def download_all_combinations(self) -> Dict[str, Any]:
        """
        下载所有symbol和indicator的排列组合数据
        
        Returns:
            下载结果统计
        """
        start_time = datetime.now()
        total_records = 0
        success_combinations = 0
        failed_combinations = 0
        
        logger.info(f"开始下载所有参数组合的基金数据，共{self.total_combinations}个组合")
        
        # 在开始下载前，清理今天的旧数据
        self.clear_today_data()
        
        for symbol in self.symbol_options:
            for indicator in self.indicator_options:
                self.current_combination += 1
                
                logger.info(f"处理组合 {self.current_combination}/{self.total_combinations}: symbol={symbol}, indicator={indicator}")
                
                try:
                    # 记录进度
                    self.log_progress(self.current_combination, self.total_combinations, total_records)
                    
                    # 使用重试机制获取数据
                    df = self.get_fund_info_with_retry(symbol=symbol, indicator=indicator)
                    
                    if df is not None and not df.empty:
                        # 解析数据
                        fund_data_list = self.parse_fund_data(df, symbol, indicator)
                        
                        # 使用批量保存
                        if self.save_fund_data_batch(fund_data_list):
                            total_records += len(fund_data_list)
                            success_combinations += 1
                            logger.info(f"组合 {self.current_combination} 处理成功，获取{len(fund_data_list)}条数据")
                        else:
                            failed_combinations += 1
                            logger.error(f"组合 {self.current_combination} 数据保存失败")
                    else:
                        logger.warning(f"组合 {self.current_combination} 未获取到数据")
                        failed_combinations += 1
                    
                    # 添加延迟，避免请求过于频繁
                    time.sleep(CONFIG['REQUEST_DELAY'])
                    
                except Exception as e:
                    failed_combinations += 1
                    logger.error(f"组合 {self.current_combination} 处理失败: {str(e)}")
                    continue
        
        end_time = datetime.now()
        duration = (end_time - start_time).total_seconds()
        
        result = {
            'start_time': start_time,
            'end_time': end_time,
            'duration_seconds': duration,
            'total_combinations': self.total_combinations,
            'success_combinations': success_combinations,
            'failed_combinations': failed_combinations,
            'total_records': total_records
        }
        
        logger.info(f"所有组合下载完成 - 耗时: {duration:.2f}秒, 成功: {success_combinations}, 失败: {failed_combinations}, 总记录数: {total_records}")
        
        return result
    
    def clear_today_data(self):
        """
        清理今天的数据，避免重复数据
        """
        try:
            today = date.today()
            deleted_count = FundNavRanking.objects.filter(nav_date=today).delete()[0]
            logger.info(f"清理今天的旧数据，删除了{deleted_count}条记录")
        except Exception as e:
            logger.error(f"清理今天数据失败: {str(e)}")
    
    def download_single_combination(self, symbol: str = "全部", indicator: str = "全部") -> Dict[str, Any]:
        """
        下载单个symbol和indicator组合的数据
        
        Args:
            symbol: symbol分类
            indicator: indicator类型
            
        Returns:
            下载结果
        """
        start_time = datetime.now()
        
        logger.info(f"开始下载单个组合数据: symbol={symbol}, indicator={indicator}")
        
        try:
            # 使用重试机制获取数据
            df = self.get_fund_info_with_retry(symbol=symbol, indicator=indicator)
            
            if df is not None and not df.empty:
                # 解析数据
                fund_data_list = self.parse_fund_data(df, symbol, indicator)
                
                # 使用批量保存
                if self.save_fund_data_batch(fund_data_list):
                    end_time = datetime.now()
                    duration = (end_time - start_time).total_seconds()
                    
                    result = {
                        'success': True,
                        'symbol': symbol,
                        'indicator': indicator,
                        'records_count': len(fund_data_list),
                        'duration_seconds': duration,
                        'start_time': start_time,
                        'end_time': end_time
                    }
                    
                    logger.info(f"单个组合下载成功 - symbol: {symbol}, indicator: {indicator}, 记录数: {len(fund_data_list)}, 耗时: {duration:.2f}秒")
                    return result
                else:
                    raise Exception("数据保存失败")
            else:
                raise Exception("未获取到数据")
                
        except Exception as e:
            end_time = datetime.now()
            duration = (end_time - start_time).total_seconds()
            
            result = {
                'success': False,
                'symbol': symbol,
                'indicator': indicator,
                'error': str(e),
                'duration_seconds': duration,
                'start_time': start_time,
                'end_time': end_time
            }
            
            logger.error(f"单个组合下载失败 - symbol: {symbol}, indicator: {indicator}, 错误: {str(e)}")
            return result


def main():
    """
    主函数
    """
    try:
        # 检查akshare是否安装
        try:
            import akshare as ak
            logger.info(f"AKShare版本: {ak.__version__}")
        except ImportError:
            logger.error("请先安装akshare: pip install akshare")
            return
        
        # 创建下载器实例
        downloader = AKShareFundRankingDownloader()
        
        # 下载所有组合数据
        result = downloader.download_all_combinations()
        
        # 输出结果
        print("\n=== 下载结果统计 ===")
        print(f"开始时间: {result['start_time']}")
        print(f"结束时间: {result['end_time']}")
        print(f"总耗时: {result['duration_seconds']:.2f}秒")
        print(f"总组合数: {result['total_combinations']}")
        print(f"成功组合: {result['success_combinations']}")
        print(f"失败组合: {result['failed_combinations']}")
        print(f"总记录数: {result['total_records']}")
        
    except Exception as e:
        logger.error(f"程序执行失败: {str(e)}")
        raise


if __name__ == '__main__':
    main()